import numpy as np

class DoubleIntegratorVehicle:
    """
    双积分小车模型
    状态变量: [x, y, vx, vy]
    控制输入: [ax, ay]
    """
    def __init__(self, dt=0.1):
        # 采样时间
        self.dt = dt
        
        # 系统矩阵 A
        self.A = np.array([
            [1.0, 0.0, dt, 0.0],
            [0.0, 1.0, 0.0, dt],
            [0.0, 0.0, 1.0, 0.0],
            [0.0, 0.0, 0.0, 1.0]
        ])
        
        # 控制矩阵 B
        self.B = np.array([
            [0.5*dt**2, 0.0],
            [0.0, 0.5*dt**2],
            [dt, 0.0],
            [0.0, dt]
        ])
        
        # 系统状态 [x, y, vx, vy]
        self.state = np.zeros(4, dtype=np.float64)
        
        # 打印系统信息
        print("双积分小车模型初始化:")
        print(f"采样时间: {dt}")
        print("系统矩阵 A:")
        print(self.A)
        print("控制矩阵 B:")
        print(self.B)
    
    def reset(self, initial_state=None):
        """重置系统状态"""
        if initial_state is not None:
            # 确保状态是浮点数并且形状正确
            if len(initial_state) != 4:
                raise ValueError(f"初始状态必须是4维向量，而不是 {len(initial_state)}")
            self.state = np.array(initial_state, dtype=np.float64)
        else:
            self.state = np.zeros(4, dtype=np.float64)
            
        print(f"车辆状态已重置为: {self.state}")
        return self.state.copy()
    
    def step(self, u):
        """
        执行一步状态更新
        u: 控制输入 [ax, ay]
        返回: 新状态 [x, y, vx, vy]
        """
        # 确保控制输入是numpy数组并限制其范围
        u = np.array(u, dtype=np.float64)
        
        # 添加控制限制，防止不稳定
        u_clipped = np.clip(u, -20.0, 20.0)
        if not np.array_equal(u, u_clipped):
            print(f"警告：控制输入被裁剪 {u} -> {u_clipped}")
        
        # 更新状态：x(k+1) = A*x(k) + B*u(k)
        new_state = np.dot(self.A, self.state) + np.dot(self.B, u_clipped)
        
        # 更新内部状态
        self.state = new_state.copy()
        
        # 检查状态的有效性
        if np.any(np.isnan(self.state)) or np.any(np.isinf(self.state)):
            print(f"警告：状态包含NaN或Inf值: {self.state}")
            # 重置为安全状态
            self.state = np.zeros(4, dtype=np.float64)
        
        return self.state.copy()
    
    def get_state_space_matrices(self):
        """返回状态空间矩阵，用于LQR控制器设计"""
        return self.A.copy(), self.B.copy()
    
    def get_state(self):
        """获取当前状态"""
        return self.state.copy()
        
    def __str__(self):
        return f"双积分小车 [位置=({self.state[0]:.2f}, {self.state[1]:.2f}), 速度=({self.state[2]:.2f}, {self.state[3]:.2f})]" 